import os
import time
import csv
import concurrent.futures
import torch
import gc
from tqdm import tqdm
from .path_utils import get_image_files_in_directory, get_video_files_in_directory


def get_gpu_memory_usage():
    """
    获取当前GPU内存使用情况
    """
    if torch.cuda.is_available():
        total_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3  # 转换为GB
        used_memory = torch.cuda.memory_allocated(0) / 1024**3  # 转换为GB
        return used_memory, total_memory
    else:
        return 0, 0


def calculate_optimal_workers():
    """
    根据GPU内存和系统资源计算最佳并行工作线程数
    """
    # 检查是否有GPU可用
    if torch.cuda.is_available():
        _, total_memory = get_gpu_memory_usage()
        
        # 根据GPU内存调整线程数
        if total_memory >= 16:  # 16GB以上GPU
            workers = 4
        elif total_memory >= 12:  # 12-16GB GPU
            workers = 3
        elif total_memory >= 8:  # 8-12GB GPU
            workers = 2
        elif total_memory >= 4:  # 4-8GB GPU
            workers = 1
    else:
        # CPU模式下使用更多线程
        import multiprocessing
        workers = max(1, multiprocessing.cpu_count() // 2)
    
    print(f"根据系统资源，设置并行工作线程数: {workers}")
    return workers


def process_batch(items, process_func, max_workers=None, progress_bar=True):
    """
    并行处理一批项目，优化版本
    
    Args:
        items: 待处理的项目列表
        process_func: 处理单个项目的函数，函数应该返回 (item, result, time_cost)
        max_workers: 最大工作线程数，如果为None则自动计算
        progress_bar: 是否显示进度条
        
    Returns:
        处理结果列表，每个元素为 (item, result, time_cost)
    """
    if not items:
        return []
    
    # 计算最佳线程数
    if max_workers is None:
        max_workers = calculate_optimal_workers()
    
    # 如果线程数为1，使用串行处理
    if max_workers <= 1:
        results = []
        for item in tqdm(items, desc="处理中", disable=not progress_bar):
            result = process_func(item)
            results.append(result)
        return results
    
    # 否则使用并行处理
    results = []
    # 使用ThreadPoolExecutor并设置超时，防止死锁
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        future_to_item = {
            executor.submit(process_func, item): item
            for item in items
        }
        
        # 获取结果，添加错误处理和超时机制
        futures_done = []
        try:
            if progress_bar:
                # 使用tqdm显示进度，但使用更高效的实现方式
                with tqdm(total=len(items), desc="处理中") as pbar:
                    # 逐个处理完成的future，避免一次性等待所有结果
                    for future in concurrent.futures.as_completed(future_to_item):
                        item = future_to_item[future]
                        try:
                            # 添加结果获取超时，避免单个任务阻塞整体进度
                            result = future.result(timeout=300)  # 5分钟超时
                            results.append(result)
                        except concurrent.futures.TimeoutError:
                            print(f"处理{item}超时")
                            results.append((item, None, 0))
                        except Exception as e:
                            print(f"处理{item}时出错: {e}")
                            results.append((item, None, 0))
                        pbar.update(1)
            else:
                # 不显示进度条的情况
                for future in concurrent.futures.as_completed(future_to_item):
                    item = future_to_item[future]
                    try:
                        result = future.result(timeout=300)  # 5分钟超时
                        results.append(result)
                    except concurrent.futures.TimeoutError:
                        print(f"处理{item}超时")
                        results.append((item, None, 0))
                    except Exception as e:
                        print(f"处理{item}时出错: {e}")
                        results.append((item, None, 0))
        except KeyboardInterrupt:
            # 优雅处理用户中断
            print("检测到用户中断，正在清理任务...")
            # 取消所有未完成的任务
            for future in future_to_item:
                if not future.done():
                    future.cancel()
            # 等待已提交任务完成
            for future in concurrent.futures.as_completed(future_to_item):
                pass
            print("任务清理完成")
    
    return results


def create_batch_processor(inferencer_class, deploy=None):
    """
    创建批量处理器函数
    
    Args:
        inferencer_class: 推理器类，如ImageInferencer或VideoInferencer
        deploy: 部署实例，如果为None则创建新的
        
    Returns:
        批量处理函数
    """
    def batch_process(items, prompts=None, save_result=True, max_workers=None, progress_bar=True):
        """
        批量处理图像或视频
        
        Args:
            items: 图像或视频路径列表
            prompts: 提示词列表，如果为None则全部使用默认提示词
            save_result: 是否保存结果到文件
            max_workers: 最大工作线程数，如果为None则自动计算
            progress_bar: 是否显示进度条
            
        Returns:
            处理结果列表
        """
        # 创建推理器实例
        inferencer = inferencer_class(deploy)
        
        # 确保prompts列表长度与items一致
        if prompts is None:
            prompts = [None] * len(items)
        elif len(prompts) != len(items):
            print("警告: prompts列表长度与items不一致，将使用默认提示词")
            prompts = [None] * len(items)
        
        # 创建处理函数
        def process_item(item_prompt):
            item, prompt = item_prompt
            start_time = time.time()
            result = inferencer.infer(item, prompt, save_result)
            end_time = time.time()
            time_cost = end_time - start_time
            return item, result, time_cost
        
        # 处理所有项目
        items_with_prompts = list(zip(items, prompts))
        results = process_batch(items_with_prompts, process_item, max_workers, progress_bar)
        
        return results
    
    return batch_process


def export_results_to_csv(results, output_file, field_names=None):
    """
    将处理结果导出到CSV文件
    
    Args:
        results: 处理结果列表，每个元素为 (file_path, result, time_cost)
        output_file: 输出CSV文件路径
        field_names: 字段名列表，如果为None则使用默认字段名
        
    Returns:
        是否导出成功
    """
    try:
        # 确保输出目录存在
        output_dir = os.path.dirname(output_file)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)
        
        # 使用默认字段名
        if field_names is None:
            field_names = ['filename', 'file_path', 'result', 'time_cost']
        
        # 写入CSV文件
        with open(output_file, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=field_names)
            writer.writeheader()
            
            for file_path, result, time_cost in results:
                # 检查结果是否为元组且长度为3（适应之前的batch_infer函数返回格式）
                if isinstance(result, tuple) and len(result) == 3:
                    actual_result = result[0]
                else:
                    actual_result = result
                
                # 构建行数据
                row = {
                    'filename': os.path.basename(file_path),
                    'file_path': file_path,
                    'result': actual_result,
                    'time_cost': f"{time_cost:.2f}"
                }
                
                # 只写入field_names中包含的字段
                row_to_write = {k: v for k, v in row.items() if k in field_names}
                writer.writerow(row_to_write)
        
        print(f"结果已导出到CSV文件: {output_file}")
        return True
    except Exception as e:
        print(f"导出结果到CSV时出错: {e}")
        return False


def process_directory(
    dir_path,
    inferencer_class,
    prompt=None,
    save_result=True,
    export_csv=None,
    max_workers=None,  # 默认为None，自动计算
    recursive=True,
    deploy=None,
    batch_size=None  # 默认None，自动计算
):
    """
    处理目录中的所有图像或视频文件，优化版本，添加内存管理和批处理优化
    
    Args:
        dir_path: 目录路径
        inferencer_class: 推理器类，如ImageInferencer或VideoInferencer
        prompt: 提示词模板，如果为None则使用默认提示词
        save_result: 是否保存结果到文件
        export_csv: 导出CSV文件路径，如果为None则不导出
        max_workers: 最大工作线程数，推荐使用1以避免GPU内存问题
        recursive: 是否递归搜索子目录
        deploy: 部署实例，如果为None则创建新的
        batch_size: 批处理大小，根据内存情况调整
        
    Returns:
        处理结果列表
    """
    # 检查目录是否存在
    if not os.path.exists(dir_path):
        print(f"目录不存在: {dir_path}")
        return []
    
    # 根据推理器类型获取相应的文件
    if 'Image' in inferencer_class.__name__:
        items = get_image_files_in_directory(dir_path, recursive)
        item_type = '图像'
    elif 'Video' in inferencer_class.__name__:
        items = get_video_files_in_directory(dir_path, recursive)
        item_type = '视频'
    else:
        print("未知的推理器类型")
        return []
    
    if not items:
        print(f"在目录{dir_path}中未找到{item_type}文件")
        return []
    
    print(f"找到{len(items)}个{item_type}文件待处理")
    
    # 计算最佳工作线程数
    if max_workers is None:
        max_workers = calculate_optimal_workers()
    
    # 计算最佳批处理大小
    if batch_size is None:
        if torch.cuda.is_available():
            _, total_memory = get_gpu_memory_usage()
            # 根据GPU内存调整批处理大小
            if total_memory >= 16:
                batch_size = 8
            elif total_memory >= 12:
                batch_size = 6
            elif total_memory >= 8:
                batch_size = 4
            elif total_memory >= 6:
                batch_size = 3
            else:
                batch_size = 2
        else:
            # CPU模式下使用较小的批处理大小
            batch_size = 4
    
    print(f"使用处理模式: {'并行' if max_workers > 1 else '串行'}, 线程数: {max_workers}, 批处理大小: {batch_size}")
    
    # 创建批量处理器，传递deploy实例
    batch_processor = create_batch_processor(inferencer_class, deploy)
    
    # 准备提示词列表
    if prompt is not None:
        # 如果提供了提示词模板，为每个文件生成提示词
        prompts = []
        for item in items:
            filename = os.path.basename(item)
            # 替换模板中的变量
            file_prompt = prompt.replace('{file_name}', filename)
            prompts.append(file_prompt)
    else:
        prompts = None
    
    # 开始处理
    start_time = time.time()
    
    # 优化的批处理方式，避免一次性加载所有数据
    results = []
    # 记录上次清理缓存的时间
    last_cache_clean_time = time.time()
    cache_clean_interval = 60  # 60秒清理一次缓存
    
    for i in range(0, len(items), batch_size):
        batch_items = items[i:i+batch_size]
        batch_prompts = prompts[i:i+batch_size] if prompts else None
        
        # 处理当前批次
        batch_results = batch_processor(batch_items, batch_prompts, save_result, max_workers)
        results.extend(batch_results)
        
        # 已处理文件数
        processed_count = min(i+batch_size, len(items))
        
        # 智能清理缓存：按时间间隔或处理一定数量文件后清理
        current_time = time.time()
        if (current_time - last_cache_clean_time > cache_clean_interval or 
            processed_count % (batch_size * 2) == 0):
            print(f"已处理{processed_count}/{len(items)}个文件，清理缓存...")
            torch.cuda.empty_cache() if torch.cuda.is_available() else None
            gc.collect()
            last_cache_clean_time = current_time
        else:
            print(f"已处理{processed_count}/{len(items)}个文件")
    
    end_time = time.time()
    
    # 统计成功和失败的数量
    success_count = sum(1 for _, result, _ in results if result is not None)
    fail_count = len(results) - success_count
    
    print(f"\n===== 批量处理完成 =====")
    print(f"总文件数: {len(items)}")
    print(f"成功处理: {success_count}")
    print(f"处理失败: {fail_count}")
    print(f"总耗时: {end_time - start_time:.2f}秒")
    print(f"平均耗时: {(end_time - start_time) / len(items):.2f}秒/文件")
    print(f"======================")
    
    # 导出CSV
    if export_csv:
        export_results_to_csv(results, export_csv)
    
    # 最终清理资源
    print("最终清理资源，释放内存...")
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    gc.collect()
    
    return results